# -*- coding: utf-8 -*-

import numpy as np
from scipy.optimize import linear_sum_assignment
import random
import rospy
import sys
import importlib.util
import os
import copy
# from tf.transformations import quaternion_matrix


# sys.path.append(current_dir)
# print(current_dir)

from pvesti.msg import UAVPose, UAVPosVel
from geometry_msgs.msg import Pose, Point, Twist, Vector3
from nav_msgs.msg import Odometry

# current_dir = os.path.dirname(os.path.abspath(__file__))
# module_name = "kalmanPredict"
# module_path = os.path.join(current_dir, f'{module_name}.py')
# spec = importlib.util.spec_from_file_location(module_name, module_path)
# kalman_module = importlib.util.module_from_spec(spec)
# spec.loader.exec_module(kalman_module)

# # 从模块中导入 KalmanFilterPredictor 类
# KalmanFilterPredictor = kalman_module.KalmanFilterPredictor


from kalmanPredict import KalmanFilterPredictor


#检测到的无人机个数
DETECT_UAV_NUM = 2
#识别到的无人机卡尔曼滤波器
kalman_predictors = {}
#当前无人机卡尔曼滤波器
current_uav_kalman_predictor = None 
# 两架无人机的区分度
THRESHOLD_DIS = 0.8


# vins旋转
get_first_vins = True
vins_rot_initial = np.eye(3) 

    
# vehicle_type = sys.argv[1]
# vehicle_num = int(sys.argv[2])

# UAV 信息类，用于存储 UAV 的时间戳、数量和位置
class UAVInfo:
    def __init__(self):
        self.seq = -1
        self.timestamp = None
        self.num = 0
        self.positions = []

class UAVPosVelPredict:
    def __init__(self):
        self.position = None
        self.timestamp = None
        self.kalman_predictor = None
        
        

UAVInfo_pre = UAVInfo()   # 之前的 UAV 信息
UAVInfo_cur = UAVInfo()   # 当前的 UAV 信息
UAVInfo_callback_cur = UAVInfo() # 回调函数每更新一次，他就更新
last_seq = -100           # 最新序号
vins_pos = Point()        # 保存来自vins位置的信息
vins_rot = np.eye(3)      # 保存vins旋转，初值单位矩阵

def quaternion_to_rotation_matrix(q):
    
    [q_x, q_y, q_z, q_w] = q

    # 计算旋转矩阵
    R = np.array([
        [1 - 2 * (q_y**2 + q_z**2), 2 * (q_x * q_y - q_z * q_w), 2 * (q_x * q_z + q_y * q_w)],
        [2 * (q_x * q_y + q_z * q_w), 1 - 2 * (q_x**2 + q_z**2), 2 * (q_y * q_z - q_x * q_w)],
        [2 * (q_x * q_z - q_y * q_w), 2 * (q_y * q_z + q_x * q_w), 1 - 2 * (q_x**2 + q_y**2)]
    ])
    
    return R


# 位置匹配类，用于匹配当前和之前的位置
class PositionMatcher:
    def __init__(self, count):
        self.global_uav_info = []
        self.total_count = count
        
    def match_positions(self, current_positions, current_timestamp):
        # ros获得的初始时间戳单位是ns
        # print(f"current_timestamp: {current_timestamp}")
        current_timestamp_sec = current_timestamp.to_sec()
        

        
        # 1. 数量达到应该检测到的数量时            
        if len(self.global_uav_info) == self.total_count:
            # 1.1 计算成本矩阵，匈牙利算法
            max_len = max(len(self.global_uav_info), len(current_positions))
            large_number = 100000.0
            cost_matrix = np.full((max_len, max_len), fill_value=large_number)
            for i in range(len(self.global_uav_info)):
                for j in range(len(current_positions)):
                    cost_matrix[i, j] = np.linalg.norm(np.array(self.global_uav_info[i].position) - np.array(current_positions[j]))
            
            # 使用匈牙利算法找到最小成本匹配
            row_idx, col_idx = linear_sum_assignment(cost_matrix)
            
            valid_matches = []
            # 1.2 更新最新位置和卡尔曼预测结果
            for r, c in zip(row_idx, col_idx):
                # 保证匹配的那俩不是最后勉强匹配上的那对，也就是权重为10000的那对（过滤掉权重高的填充匹配）
                if cost_matrix[r, c] < large_number:
                    valid_matches.append((r,c))
                    
                    last_position = self.global_uav_info[r].position
                    last_timestamp = self.global_uav_info[r].timestamp
                    dt = current_timestamp_sec - last_timestamp
                    
                    self.global_uav_info[r].position = current_positions[c]
                    self.global_uav_info[r].timestamp = current_timestamp_sec
                    self.global_uav_info[r].kalman_predictor.predict(current_positions[c], dt)
                    
                else:
                    # print(f"Match bewteen {r} and {c} is invalid!")
                    pass
        
        # 2. 数量不够总数时        
        elif len(self.global_uav_info) < self.total_count:
            # 2.1.在self.global_uav_info加入新的无人机
            
            # 2.1.1 原本self.global_uav_info为空，直接加入
            if len(self.global_uav_info) == 0:
                for idx in range(len(current_positions)):
                    uavInfo = UAVPosVelPredict()
                    uavInfo.position = current_positions[idx]
                    uavInfo.timestamp = current_timestamp_sec
                    # 加入卡尔曼预测
                    uavInfo.kalman_predictor = KalmanFilterPredictor()
                    uavInfo.kalman_predictor.initialize_state(current_positions[idx])
                    
                    self.global_uav_info.append(uavInfo)
                    
                
            # 2.1.2 判断current_positions中是否有新的无人机位置，并找到其在current_positions中的索引idx，加入self.global_uav_info
            else:
                for pos in current_positions:
                    distances = np.array([np.linalg.norm(np.array(pos)-np.array(global_pos.position)) for global_pos in self.global_uav_info])
                    if len(distances) > 0:
                        if np.min(distances) < THRESHOLD_DIS:
                            # 这个pos在原有pos中
                            pass
                        else:
                            if len(self.global_uav_info) < self.total_count:
                                uavInfo = UAVPosVelPredict()
                                uavInfo.position = pos
                                uavInfo.timestamp = current_timestamp_sec
                                # 加入卡尔曼预测
                                uavInfo.kalman_predictor = KalmanFilterPredictor()
                                uavInfo.kalman_predictor.initialize_state(pos)
                                
                                self.global_uav_info.append(uavInfo)
            
            # 2.2 更新self.global_uav_info并进行卡尔曼预测
            max_len = max(len(self.global_uav_info), len(current_positions))
            large_number = 100000.0
            cost_matrix = np.full((max_len, max_len), fill_value=large_number)
            for i in range(len(self.global_uav_info)):
                for j in range(len(current_positions)):
                    cost_matrix[i, j] = np.linalg.norm(np.array(self.global_uav_info[i].position) - np.array(current_positions[j]))
            
            # 使用匈牙利算法找到最小成本匹配
            row_idx, col_idx = linear_sum_assignment(cost_matrix)
            
            valid_matches = []
            # 过滤掉权重高的填充匹配
            for r, c in zip(row_idx, col_idx):
                if cost_matrix[r, c] < large_number:
                    valid_matches.append((r,c))
                    
                    last_timestamp = self.global_uav_info[r].timestamp
                    dt = current_timestamp_sec - last_timestamp
                    
                    self.global_uav_info[r].position = current_positions[c]
                    self.global_uav_info[r].timestamp = current_timestamp_sec
                    self.global_uav_info[r].kalman_predictor.predict(current_positions[c], dt)
                    
                else:
                    # print(f"Match bewteen {r} and {c} is invalid!")
                    pass
        
        
        # 检测并删除长时间未更新的位置
        TIME_THRES = 2
        # 倒序遍历删除，否则会越界，ps:range的遍历是不会包含最后一个数得到，所以终止值是-1
        for i in range(len(self.global_uav_info)-1, -1, -1):
            if abs(self.global_uav_info[i].timestamp - current_timestamp_sec) > TIME_THRES:
                del self.global_uav_info[i]
        

                
                        
matcher = PositionMatcher(DETECT_UAV_NUM)           
        

# UAV 相对位置信息回调函数
def UAV_pose_callback(data):
    global UAVInfo_callback_cur, UAVInfo_pre, UAVInfo_cur, vins_pos, vins_rot
    
    UAVInfo_callback_cur.seq = data.header.seq
    UAVInfo_callback_cur.timestamp = data.header.stamp
    UAVInfo_callback_cur.num = data.uav_count
        
    # rospy.loginfo("UAV_pose_callback':")
    # rospy.loginfo("UAVInfo_pre seq and positions':{}, {}".format(UAVInfo_pre.seq, UAVInfo_pre.positions))
    # rospy.loginfo("UAVInfo_cur seq and positions':{}, {}".format(UAVInfo_cur.seq, UAVInfo_cur.positions))
    
    # # 若没有更新则清空（！！！但是，没有更新值，这个地方就不掉用，所以更新操作写在这里没有用！！！）
    # if UAVInfo_pre.seq > 0 and (UAVInfo_pre.seq == data.header.seq):
    #     rospy.loginfo("No new data received. Clearing current UAV info.")
    #     UAVInfo_cur.seq = -1
    #     UAVInfo_cur.timestamp = None
    #     UAVInfo_cur.num = 0
    #     UAVInfo_cur.positions = []
    # # 若更新了，将UAVInfo_cur的positions加入vins的绝对坐标
    # else:
    
    # 一定要先清空
    UAVInfo_callback_cur.positions = []
    
    # 没有订阅到vins时，vins_pos直接是（0，0，0）
    vins_pos_arr = np.array([vins_pos.x, vins_pos.y, vins_pos.z])
    pose_absolute_rots = []
    for pose_relative in data.uav_positions:
        # 坐标转换
        pose_relative_arr = np.array([pose_relative.position.x, pose_relative.position.y, pose_relative.position.z])
        pose_absolute_arr_rot = np.dot(vins_rot, pose_relative_arr) + vins_pos_arr
        # pose_absolute_arr = pose_relative_arr + vins_pos_arr
        
        # 转换为普通元组
        pose_absolute_rot = tuple(pose_absolute_arr_rot)
        # pose_absolute = tuple(pose_absolute_arr)
        UAVInfo_callback_cur.positions.append(pose_absolute_rot)
        
    
    # 在回调函数里匹配，虽然频率比较合适，但无法避免lidar识别算法识别不到时，话题不发布新的停顿在那里的情况
    # positions_cur = UAVInfo_callback_cur.positions
    # cur_time_ns = rospy.get_rostime()
    # matcher.match_positions(positions_cur, cur_time_ns)
    # uav_postions_global = matcher.global_uav_info
    
    # UAVInfo_pre = copy.deepcopy(UAVInfo_cur)
    
    
    # rospy.loginfo("position with Vins's rotation': {}".format(pose_absolute_rots))
    # rospy.loginfo("position without Vins's rotation': {}\n".format(UAVInfo_cur.positions))
        

# UAV 绝对位置信息回调函数
def vins_callback(data):
    global vins_pos, vins_rot, vins_rot_initial, get_first_vins
    
    # 处理vins初值旋转不是单位矩阵的情况
    if get_first_vins:
        vins_quat = [data.pose.pose.orientation.x, data.pose.pose.orientation.y, 
                 data.pose.pose.orientation.z, data.pose.pose.orientation.w]
        vins_rot = quaternion_to_rotation_matrix(vins_quat)
        # 求当前旋转矩阵的逆
        vins_rot_initial = np.linalg.inv(vins_rot)
        
        get_first_vins = False
    
    vins_pos=Point()
    vins_pos=data.pose.pose.position
    
    vins_quat = [data.pose.pose.orientation.x, data.pose.pose.orientation.y, 
                 data.pose.pose.orientation.z, data.pose.pose.orientation.w]
    vins_rot = quaternion_to_rotation_matrix(vins_quat)
    
    vins_rot = np.dot(vins_rot, vins_rot_initial)
    
    # rospy.loginfo("Vins's rotation:\n {}".format(vins_rot))


# 主函数，用于测试位置匹配和卡尔曼滤波
def main():
    matcher = PositionMatcher(DETECT_UAV_NUM)

    dt = 1.0
    kalman_predictors = {}

    # # 测试数据
    # m = 5  # 时间步数
    # n = 3  # 物体数量
    # positions = [[(1,1,2), (3,3,2), (6,1,2)],
    #              [(3.2, 3.2, 2), (1.7, 1.8, 2), (6.1, 1.8, 2)],
    #              [(6.4, 1.7, 2), (1.2, 1.2, 2), (3.9, 3.9, 2)],
    #              [(1.5, 1.5, 2), (3.3, 3.4, 2), (6.7, 1.2, 2)],
    #              [(1.3, 1.7, 2), (6.8, 1.6, 2), (3.2, 3.9, 2)]]
    
    m = 5  # 时间步数
    n = 3
    positions = [
        [(6.1, 1.3, 2.1), (2.5, 3.1, 2.1)],
        [(1, 1, 2), (3, 3, 2), (6, 1, 2)],  
        [(3.2, 3.2, 2), (1.7, 1.8, 2)], 
        [(1.2, 1.2, 2), (3.9, 3.9, 2), (6.4, 1.7, 2), (8.5, 0.5, 2)],  
        [(1.5, 1.5, 2), (3.3, 3.4, 2), (6.7, 1.2, 2)], 
        [(1.3, 1.7, 2), (3.2, 3.9, 2)]]  
    

    last_match_uav_size = 0
    
    for i in range(m):
        positions_cur = positions[i]

        # 匹配位置
        matcher.match_positions(positions_cur)
        uav_postions_global = matcher.global_uav_info
        
        if i == 0:
            # 初始化卡尔曼滤波器
            for idx in range(len(uav_postions_global)):
                kalman_predictors[idx] = KalmanFilterPredictor(dt)
                kalman_predictors[idx].initialize_state(uav_postions_global[idx].position)
        else:
            # 初始化新识别的uav对应的卡尔曼滤波器
            if len(uav_postions_global) > last_match_uav_size:
                for idx in range(last_match_uav_size, len(uav_postions_global)):
                    # print(idx)
                    kalman_predictors[idx] = KalmanFilterPredictor(dt)
                    kalman_predictors[idx].initialize_state(uav_postions_global[idx].position)

        # 卡尔曼预测，没有识别到的物体的速度和位置沿用上一时刻的速度和位置         
        for idx in range(len(uav_postions_global)):
            kalman_predictors[idx].predict(uav_postions_global[idx].position, dt)
        
        last_match_uav_size = len(uav_postions_global)
        
        # 打印实时的滤波位置和速度
        print("Time {}: \n".format(i))
        for j in range(len(uav_postions_global)):
            detectied_positions = uav_postions_global[j]
            cal_velocities = kalman_predictors[j].filtered_velocities[-1]
            print("Object {} position: {}\n".format(j, detectied_positions))
            print("Object {} velocitiy: {}\n".format(j, cal_velocities))
            
            



# 主函数2，用于 ROS 节点，接收 UAV 位置信息并发布位置信息和速度
def main2():
    rospy.init_node('pos_vel_transfer')
    rospy.Subscriber('/vins_estimator/odometry', Odometry,vins_callback, queue_size=1)#从vins中得到中心无人机的绝对位置信息
    # rospy.Subscriber(vehicle_type+'_'+str(i)+'/vins_estimator/odometry', Odometry,vins_callback, queue_size=1)#从vins中得到中心无人机的绝对位置信息
    rospy.Subscriber('/uav_positions', UAVPose, UAV_pose_callback, queue_size=1)#从雷达发布节点处得到无人机数量、相对位置信息uav_count、geometry_msgs/Pose
    pub = rospy.Publisher('/uav_pos_vel', UAVPosVel, queue_size=10)#给处理避障的节点发布无人机数量、位置、速度信息uav_count、geometry_msgs/Pose、geometry_msgs/Twist
    

    
    
    
    global UAVInfo_callback_cur, UAVInfo_pre, UAVInfo_cur, vins_pos, last_seq
    
    first_run = True
    ROS_RATE = 1
    rate = rospy.Rate(ROS_RATE)
    
    last_match_uav_size = 0
    
    # 卡尔曼的估计频率就是dt，但是是预测速度+
    while not rospy.is_shutdown():
        
        # 一定要用深拷贝
        UAVInfo_cur = copy.deepcopy(UAVInfo_callback_cur)
        
        # rospy.loginfo("UAVInfo_pre seq and positions':{}, {}".format(UAVInfo_pre.seq, UAVInfo_pre.positions))
        # rospy.loginfo("UAVInfo_cur seq and positions':{}, {}".format(UAVInfo_cur.seq, UAVInfo_cur.positions))
        
        # # 第一次清空
        # if UAVInfo_cur.seq >= 0 and UAVInfo_pre.seq >= 0 and (UAVInfo_cur.seq == UAVInfo_pre.seq):
        #     rospy.loginfo("No new data received. Clearing current UAV info.")
        #     UAVInfo_cur.seq = -1
        #     UAVInfo_cur.timestamp = None
        #     UAVInfo_cur.num = 0
        #     UAVInfo_cur.positions = []
        
        if last_seq == UAVInfo_callback_cur.seq:
            rospy.loginfo("last_seq = UAVInfo_callback_cur.seq, Clearing current UAV info.")
            UAVInfo_cur.seq = -1
            UAVInfo_cur.timestamp = None
            UAVInfo_cur.num = 0
            UAVInfo_cur.positions = []
            
        dt = 1.0/ROS_RATE
        if UAVInfo_cur.timestamp is not None and UAVInfo_pre.timestamp is not None:
            dt = UAVInfo_cur.timestamp.to_sec() - UAVInfo_pre.timestamp.to_sec()
            
        
        # rospy.loginfo("UAVInfo_cur seq and positions':{}, {}".format(UAVInfo_cur.seq, UAVInfo_cur.positions))
        positions_cur = UAVInfo_cur.positions
        cur_time_ns = rospy.get_rostime()
        matcher.match_positions(positions_cur, cur_time_ns)
        uav_postions_global = matcher.global_uav_info
        
        
        if first_run:
            # 初始化当前uav的卡尔曼滤波器 
            current_uav_kalman_predictor = KalmanFilterPredictor(dt)
            current_uav_kalman_predictor.initialize_state([vins_pos.x, vins_pos.y, vins_pos.z])
            first_run = False
            

        
        # 当前uav卡尔曼滤波器-预测
        current_position = [vins_pos.x, vins_pos.y, vins_pos.z]
        current_uav_kalman_predictor.predict(current_position, dt)
            
        msg = UAVPosVel()
        msg.header.stamp = UAVInfo_cur.timestamp
        msg.uav_count = len(uav_postions_global)
        
        # 增加当前无人机的位置和速度
        pose_cur = Pose()
        pose_cur.position = Point(*current_uav_kalman_predictor.filtered_positions[-1])
        twist_cur = Twist()
        twist_cur.linear = Vector3(*current_uav_kalman_predictor.filtered_velocities[-1])
        msg.uav_positions.append(pose_cur)
        msg.uav_velocities.append(twist_cur)
        
        pose_cur_info = [pose_cur.position.x, pose_cur.position.y, pose_cur.position.z]
        vel_cur_info = [twist_cur.linear.x, twist_cur.linear.y, twist_cur.linear.z]
        rospy.loginfo("current UAV position': {}".format(pose_cur_info))
        rospy.loginfo("current UAV velocity': {}".format(vel_cur_info))
        
        
        for i in range(len(uav_postions_global)):
            pose = Pose()
            pose.position = Point(*uav_postions_global[i].kalman_predictor.filtered_positions[-1])
            twist = Twist()
            twist.linear = Vector3(*uav_postions_global[i].kalman_predictor.filtered_velocities[-1])
            msg.uav_positions.append(pose)
            msg.uav_velocities.append(twist)
            
            pose_info = [pose.position.x, pose.position.y, pose.position.z]
            vel_info = [twist.linear.x, twist.linear.y, twist.linear.z]
            
            
            rospy.loginfo("UAV {} position': {}".format(i, pose_info))
            rospy.loginfo("UAV {} velocity': {}".format(i, vel_info))
            
        # 发布消息    
        pub.publish(msg)
        
        # 一定要用深拷贝
        UAVInfo_pre = copy.deepcopy(UAVInfo_cur)
        if UAVInfo_pre.seq != -1:
            last_seq = UAVInfo_pre.seq
            
        rate.sleep()

if __name__ == "__main__":
    main2()